import os
import librosa
import numpy as np
import torch
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from torch.utils.data import Dataset, DataLoader, TensorDataset


def dataloader(DataDir, N_MFCC, MAX_PAD_LEN, BATCH_SIZE):
    """实例化MakeDataset类"""
    loader = MakeDataset(DataDir, N_MFCC, MAX_PAD_LEN, BATCH_SIZE)
    return loader


class MakeDataset(torch.utils.data.Dataset):
    def __init__(self, DataDir, N_MFCC, MAX_PAD_LEN, BATCH_SIZE):
        self.DataDir = DataDir  # 数据所在目录
        self.N_MFCC = N_MFCC  # 提取mfcc特征数量
        self.MAX_PAD_LEN = MAX_PAD_LEN
        self.BATCH_SIZE = BATCH_SIZE

    def load_data(self):
        """
        加载数据集，返回文件列表、文件标签（0起计）
        :param DataDir:数据集目录路径
        """
        # 找到所有的.wav音频文件
        f_wavs = []
        labels = []
        for root, dirs, files in os.walk(self.DataDir):
            for fname in files:
                if fname.endswith('.wav'):
                    f_wavs.append(os.path.join(root, fname))  # 获取文件路径
                    labels.append(int(fname.split('-')[2]) - 1)  # 获取文件标签，并且减去1，因为标签是从0开始的
        # print(f'Found .wav count: {len(f_wavs)}')
        # print(f'label classes count: {len(set(labels))}')
        return f_wavs, labels

    def extract_features(self, file_path, sr=22050):
        """
        提取音频特征
        :param file_path: 音频文件路径
        :param sr: 采样率
        :param max_pad_len: 最大MFCC特征系数数量
        """
        # 加载音频文件
        audio, _ = librosa.load(file_path, sr=sr)
        # 提取mfcc特征
        mfcc = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=self.N_MFCC)
        # print(mfcc.shape, type(mfcc))  # (13, 143)
        # 填充或截断mfcc特征
        if mfcc.shape[1] > self.MAX_PAD_LEN:
            cut_head = int((mfcc.shape[1] - self.MAX_PAD_LEN) / 2)
            mfcc = mfcc[:, cut_head:cut_head + self.MAX_PAD_LEN]
        else:
            mfcc = mfcc[:, :self.MAX_PAD_LEN]
            mfcc = np.pad(mfcc, ((0, 0), (0, self.MAX_PAD_LEN - mfcc.shape[1])), mode='constant', constant_values=0)
        # 转置前，形状是(n_mfcc, time_steps).
        # 转置原因： LSTM 网络期待的输入形状是(batch_size, sequence_length, input_dim)
        return mfcc.T

    def load(self):
        """
        加载数据。
        :param DataDir: 数据所在目录
        :param N_MFCC: 音频文件要提取的特征数量
        """
        f_wavs, y = self.load_data()  # 获取文件路径和标签
        X = [self.extract_features(f) for f in f_wavs]  # 提取mfcc特征
        features = np.array(X)
        label = np.array(y)

        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(features, label, test_size=0.3, random_state=42,
                                                            shuffle=True)
        X_test, X_val, y_test, y_val = train_test_split(X_test, y_test, test_size=0.5, random_state=42, shuffle=True)

        # # 数据标准化
        # scaler = StandardScaler()
        # X_train = scaler.fit_transform(X_train)
        # X_test = scaler.transform(X_test)
        # X_val = scaler.transform(X_val)

        # 转换为PyTorch张量
        X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
        y_train_tensor = torch.tensor(y_train, dtype=torch.long)
        trian_dataset = TensorDataset(X_train_tensor, y_train_tensor)
        train_loader = DataLoader(trian_dataset, batch_size=self.BATCH_SIZE, shuffle=True)

        X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
        y_test_tensor = torch.tensor(y_test, dtype=torch.long)
        test_dataset = TensorDataset(X_test_tensor, y_test_tensor)
        test_loader = DataLoader(test_dataset, batch_size=self.BATCH_SIZE, shuffle=True)

        X_val_tensor = torch.tensor(X_val, dtype=torch.float32)
        y_val_tensor = torch.tensor(y_val, dtype=torch.long)
        val_dataset = TensorDataset(X_val_tensor, y_val_tensor)
        val_loader = DataLoader(val_dataset, batch_size=self.BATCH_SIZE, shuffle=True)

        return train_loader, test_loader, val_loader
